package engine

import (
	"sync"
	"time"
)

type RateLimiter struct {
	rate       int       // 每秒允许的请求数
	burst      int       // 突发请求数量
	tokens     int       // 当前可用令牌数
	lastUpdate time.Time // 上次更新时间
	mu         sync.Mutex
}

func NewRateLimiter(rate int) *RateLimiter {
	return &RateLimiter{
		rate:       rate,
		burst:      rate * 2, // 突发流量设为正常速率的2倍
		tokens:     rate,
		lastUpdate: time.Now(),
	}
}

func (rl *RateLimiter) Allow() bool {
	rl.mu.Lock()
	defer rl.mu.Unlock()

	now := time.Now()
	elapsed := now.Sub(rl.lastUpdate)

	// 根据时间流逝补充令牌
	newTokens := int(elapsed.Seconds() * float64(rl.rate))
	if newTokens > 0 {
		rl.tokens += newTokens
		if rl.tokens > rl.burst {
			rl.tokens = rl.burst
		}
		rl.lastUpdate = now
	}

	// 尝试获取令牌
	if rl.tokens > 0 {
		rl.tokens--
		return true
	}

	return false
}

func (rl *RateLimiter) SetRate(rate int) {
	rl.mu.Lock()
	defer rl.mu.Unlock()

	rl.rate = rate
	rl.burst = rate * 2
}

func (rl *RateLimiter) GetRate() int {
	rl.mu.Lock()
	defer rl.mu.Unlock()
	return rl.rate
}
